# This implements a speed limited socket in Python.
# The speed limits are imposed on an application level, 
# i.e., if the download limit is 1024 kb/s then _only_ 1024 kb/s
# can be downloaded through all lsockets in the _entire_ 
# application.

from __future__ import with_statement
from threading import Condition, Thread, Lock
from time import sleep
from select import select
import socket

class Timeout(Exception):
    def __init__(self, msg):
        super(Timeout, self).__init__(msg)
        
class LSocket(object):
    # Class variables.
    TIMEOUT = 0.1
    RETRY_COUNT = 100
    UPLOAD_LIMIT = 1024 # measured in kb/s
    DOWNLOAD_LIMIT = 1024 # measured in kb/s

    write_access_lock = Condition()
    bytes_sent = 0
    read_access_lock = Condition()
    bytes_received = 0

    # Class methods.
    @classmethod
    def set_upload_limit(cls, limit):
        cls.UPLOAD_LIMIT = limit

    @classmethod
    def set_download_limit(cls, limit):
        cls.DOWNLOAD_LIMIT = limit

    @classmethod
    def reset_counters(cls):
        with cls.write_access_lock:
            cls.bytes_sent = 0
            cls.write_access_lock.notifyAll()
        with cls.read_access_lock:
            cls.bytes_received = 0
            cls.read_access_lock.notifyAll()
        
    @classmethod
    def _timed_send(cls, _socket, data, timeout):
        _, wr, _ = select([], [_socket], [], timeout)
        if wr != []:
            return _socket.send(data)
        else:
            raise Timeout('Timed out waiting to send data.')

    @classmethod
    def _send(cls, _socket, data):
        with cls.write_access_lock:
            # Wait until we may write again.
            while cls.UPLOAD_LIMIT * 1024 - cls.bytes_sent == 0:
                cls.write_access_lock.wait()
            # Write the data to the actual socket.
            bytes_left = cls.UPLOAD_LIMIT * 1024 - cls.bytes_sent
            if len(data) > bytes_left:
                # only write the allowed amount of data.
                data = data[:bytes_left]
            bytes_sent = cls._timed_send(_socket, data, cls.TIMEOUT)
            cls.bytes_sent += bytes_sent
            return bytes_sent

    @classmethod
    def _timed_recv(cls, _socket, bufsize, timeout):
        rd, _, _ = select([_socket], [], [], timeout)
        if rd != []:
            return _socket.recv(bufsize)
        else:
            raise Timeout('Timed out waiting for data to arrive on socket.')

    @classmethod
    def _recv(cls, _socket, bufsize):
        with cls.read_access_lock:
            # Wait until allowed to receive again.
            while cls.DOWNLOAD_LIMIT * 1024 - cls.bytes_received == 0:
                cls.read_access_lock.wait()
            # Read data from the socket.
            bytes_left = cls.DOWNLOAD_LIMIT * 1024 - cls.bytes_received
            if bytes_left < bufsize:
                # Only read the allowed amount of data.
                bufsize = bytes_left
            data = cls._timed_recv(_socket, bufsize, cls.TIMEOUT)
            cls.bytes_received += len(data)
            return data

    def __init__(self, sock = None):
        super(LSocket, self).__init__()
        if sock:
            self.sock = sock
        else:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        #TCP_KEEP_ALIVE forhindre socket i at lukkes pga inaktivitet.
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)


    def send(self, data):
        retries = 0
        while retries < LSocket.RETRY_COUNT:
            try:
                return LSocket._send(self.sock, data)
            except Timeout:
                retries += 1
        raise Timeout('Timed out trying to send data.')

    def recv(self, bufsize):
        retries = 0
        while retries < LSocket.RETRY_COUNT:
            try:
                return LSocket._recv(self.sock, bufsize)
            except Timeout:
                retries += 1
        raise Timeout('Timed out trying to receive data.')

    def connect(self, address):
        self.sock.connect(address)
    
    def bind(self, address):
        self.sock.bind(address)

    def accept(self):
        conn, addr = self.sock.accept()
        new_conn = LSocket(conn)
        return (new_conn, addr)

    def listen(self, backlog):
        self.sock.listen(backlog)

    def close(self):
        self.sock.close()
    
    def settimeout(self, timeout):
        self.sock.settimeout(timeout)
    
    def getpeername(self):
        return self.sock.getpeername()

class LSocketMaintenance(Thread):
    INSTANCE = None

    def __init__(self):
        if LSocketMaintenance.INSTANCE != None:
            raise Exception('Singleton error.')
        
        super(LSocketMaintenance, self).__init__()
        self._shutdown = False
        self._shutdown_signal = Lock()
        LSocketMaintenance.INSTANCE = self

    @classmethod
    def shutdown(cls):
        cls.INSTANCE._shutdown = True
        cls.INSTANCE._shutdown_signal.acquire()
        cls.INSTANCE._shutdown_signal.release()

    def run(self):
        self._shutdown_signal.acquire()
        
        while not self._shutdown:
            # Reset counters.
            LSocket.reset_counters()
            # Sleep for another second.
            sleep(1)

        self._shutdown_signal.release()
        self._shutdown = False

_lsocket_maint = LSocketMaintenance()
_lsocket_maint.start()
